"""
A multi-symbols bt demo
@Date: 2022/3/17
@Author: pyn
"""
import sqlite3
import pandas as pd
from wtpy import WTSBarStruct
import time
from wtpy.ExtModuleDefs import BaseExtDataLoader
from gym.spaces import Space, Box, Discrete, Dict, MultiDiscrete
from wtpy.WtBtEngine import WtBtEngine, EngineType
from wtpy.StrategyDefs import BaseCtaStrategy, CtaContext, BaseHftStrategy, HftContext
from wtpy.ExtModuleDefs import BaseExtDataLoader
from typing import Union, Type
from wtpy import WtExecApi, BaseExtExecuter


def sqlite2df(db_addr, table_name):
    conn = sqlite3.connect(db_addr)
    return pd.read_sql("select * from %s" % (table_name), conn)


class MyDataloader_srif(BaseExtDataLoader):
    def load_final_his_bars(self, stdCode: str, period: str, feeder) -> bool:
        """"""
        print("enter MyDataloader...")
        params2tbname = {
            'IFd1': 'ifd', 'IFm1': 'if1m', 'SRd1': 'srd', 'SRm1': 'sr1m'
        }
        df = sqlite2df('./config/srif.db', params2tbname.get(stdCode.split('.')[1] + period))

        def process_and_feeder(df, **kwargs):
            df['Unnamed: 0'] = df['Unnamed: 0'].str[0:19]
            df['date'] = df.apply(lambda x: x['Unnamed: 0'][0:10], axis=1)
            df['date'] = df['date'].str.replace('-', '/')
            if 'time_value' in kwargs:
                df['time'] = kwargs['time_value']
            else:
                df['time'] = df.apply(lambda x: x['Unnamed: 0'][-8:], axis=1)

            df = df.rename(columns={
                'date': 'date', 'time': 'time', 'open': 'open', 'high': 'high', 'low': 'low', 'close': 'close',
                'volume': 'vol'
            })
            df['date'] = df['date'].astype('datetime64').dt.strftime('%Y%m%d').astype('int64')
            df['time'] = (df['date'] - 19900000) * 10000 + df['time'].str.replace(':', '').str[:-2].astype('int')
            BUFFER = WTSBarStruct * len(df)
            buffer = BUFFER()

            def assign(procession, buffer):
                tuple(map(lambda x: setattr(buffer[x[0]], procession.name, x[1]), enumerate(procession)))

            df.apply(assign, buffer=buffer)
            ff=feeder(buffer, len(df))

            print(ff)

        if period == 'd1':
            process_and_feeder(df, time_value='00:00:00')
        elif period == 'm1':
            process_and_feeder(df)

        return True


class MyStra(BaseCtaStrategy):
    """
    1\设置为主K的订阅行情会决定oncalc,oncalcdone执行多少次，并且策略scheduled多少次。
    2\只订阅一个行情时，onbar和oncalc数量也不同。时间戳有异常如下，经过处理后，44分有两行触发了两次onbar
    2020-12-30 22:44:00
    2020-12-30 22:44:59.999999
    换了时间正常的srd数据，onbar,oncalc/done次数完全一致。

    """

    def __int__(self, name):
        """"""
        self.name = name

    def on_init(self, context: CtaContext):
        """"""
        self.info = {}
        self.info['onbarcode'] = []
        self.info['onbarperiod'] = []
        self.on_bar_num = 0
        self.on_calc_num = 0
        self.on_calc_done_num = 0
        context.stra_log_text("====>> on_init")
        # context.stra_get_bars('CFFEX.IF.HOT', 'd1', 50, isMain=False)
        # context.stra_get_bars('CFFEX.IF.HOT', 'm1', 50, isMain=True)
        context.stra_get_bars('CZCE.SR.HOT', 'd1', 50, isMain=True)
        # context.stra_get_bars('CZCE.SR.HOT', 'm1', 50, isMain=True)
        print("tdate:%s date:%s time:%s" % (
            context.stra_get_tdate(), context.stra_get_date(), context.stra_get_time()))

        # TODO 在这里是否能取得底层已转化好的date/time

        bartimes=context.stra_prepare_bars('CZCE.SR.HOT', 'd1', 50).bartimes
        print(bartimes)


    def on_bar(self, context: CtaContext, stdCode: str, period: str, newBar: dict):
        """"""
        self.on_bar_num += 1
        self.info['onbarcode'].append(stdCode)
        self.info['onbarperiod'].append(period)
        context.stra_log_text("====>> on_bar")
        print("tdate:%s date:%s time:%s stdCode:%s period:%s" % (
            context.stra_get_tdate(), context.stra_get_date(), context.stra_get_time(), stdCode, period))
        print('newBar', newBar)

    def on_calculate(self, context: CtaContext):
        """"""
        self.on_calc_num += 1
        context.stra_log_text("====>> on_calculate")
        print("tdate%s date:%s time:%s" %
              (context.stra_get_tdate(), context.stra_get_date(), context.stra_get_time()))
        bartimes = context.stra_get_bars('CZCE.SR.HOT', 'd1', 50).bartimes
        print(bartimes)
    def on_calculate_done(self, context: CtaContext):
        """"""
        self.on_calc_done_num += 1
        context.stra_log_text("====>> on_calculate_done")
        print("tdate%s date:%s time:%s" %
              (context.stra_get_tdate(), context.stra_get_date(), context.stra_get_time()))

    def on_backtest_end(self, context: CtaContext):
        print('on bar num', self.on_bar_num)
        print('on calc num', self.on_calc_num)
        print('on calc done num', self.on_calc_done_num)
        # print(self.info['onbarcode'])
        # print(self.info['onbarperiod'])


if __name__ == '__main__':
    t1 = time.time()
    print('t1', t1)
    engine = WtBtEngine(EngineType.ET_CTA, logCfg='./config/research/logcfgbt.yaml')
    engine.set_extended_data_loader(loader=MyDataloader_srif(), bAutoTrans=True)
    engine.init('.\\config\\common\\', '.\\config\\research\\configbt_cta.yaml')
    engine.configBacktest(201001031500, 201112311500)
    engine.configBTStorage(mode='csv', path='./storage')
    engine.commitBTConfig()

    stra = MyStra('s1')
    engine.set_cta_strategy(strategy=stra, hook=True)
    engine.run_backtest(bAsync=False)
    print('t2-t1 %.4f秒' % (time.time() - t1))
    kw = input('any key to exit\n')
    engine.release_backtest()
